#ifndef PCPS_MOVE_HEADER
#define PCPS_MOVE_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/coeff_value.h>
#include <src/random.h>
#include <src/arena.h>

namespace pcps {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief   checks whether the initial configuration's wavefunction value has a large enough magnitude
  ///
  /// \param[in]     userinp  user input options
  /// \param[in]     cvf      object used to compute wavefunction values
  /// \param[in]     config   the lattice configuration
  ///
  /// \return true if the magnitude is large enough, false otherwise
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template<class SCALAR, class PART, class REF> bool good_init_config_value(const pcps::Input & userinp,
                                                                            pcps::CoeffValueFunc<SCALAR, PART, REF> & cvf,
                                                                            const int * const config)
  {

    // get the wavefunction's value
    cvf.init(config);
    const pcps::MantExp<SCALAR> v = cvf.get_wfn_value(config);

    // if the value is exactly zero, the configuration is unacceptable
    if ( v.mantissa() == pcps::zero<SCALAR>() ) return false;

    // if the wavefunction's magnitude is too small, the configuration is unacceptable
    if ( double(v.exponent()) / std::sqrt(double(userinp.nparticles())) < userinp.init_config_thresh() ) {

      //std::printf("rejected a configuration with an exponent of %i\n", v.exponent());
      //std::cout << boost::format("rejected a configuration with an exponent of %i") % v.exponent() << std::endl;

      return false;
    
    }

    std::printf("chose a configuration with coefficient = %s x 10^(%i)\n", pcps::format_number("%8.4f", v.mantissa()).c_str(),
                                                                           v.exponent());
    //std::cout << boost::format("chose a configuration with coefficient = %s x 10^(%i)")
    //             % pcps::format_number("%8.4f", v.mantissa())
    //             % v.exponent()
    //          << std::endl;

    // otherwise, the configuration is acceptable
    return true;

  }

  //-------------------------------------------------------------------------------
  // pcps::MoveFunc -- A function object that performs random moves for use in
  //                   Markov-Chain Monte-Carlo.
  //
  //                   This object is templated on the scalar (SCALAR),
  //                   particle (PART), and reference (REF) types.
  //
  //                   The default of the template simply raises an error
  //                   to tell the user their choice of SCALAR, PART, and REF
  //                   was not valid.
  //
  //                   The various specializations of the template define
  //                   the actual behavior.
  //-------------------------------------------------------------------------------

  template <class SCALAR, class PART, class REF> class MoveFunc {

    public:

    // constructor
    MoveFunc(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, const pcps::Wavefunction<SCALAR> & wfn, pcps::Arena & arena) {
      this->throw_error();
    }

    // function to set the initial configuration
    template<class GENERATOR> void initialize_configuration(const pcps::Input & userinp,
                                                            pcps::CoeffValueFunc<SCALAR, pcps::Part_Spin, REF> & cvf,
                                                            int * const config,
                                                            GENERATOR & gen)
    {
      this->throw_error();
    }

    // function to propose and make a move and return the ratio of proposal densities for the configs
    template <class GENERATOR> double propose(int & n_affected, int * const affected, int * const config, GENERATOR & gen) { this->throw_error(); return 1; }

    // function to update internal data after a move has been accepted
    void accept(const int n_affected, const int * const affected, const int * const new_config) { this->throw_error(); }

    // function to tell the user that their choice of template parameters was invalid
    static void throw_error() {
      throw pcps::Exception( (boost::format("no implementation for pcps::MoveFunc with template parameters SCALAR = %s, PART = %s, REF = %s")
                               % pcps::type_lookup(typeid(SCALAR))
                               % pcps::type_lookup(typeid(PART))
                               % pcps::type_lookup(typeid(REF))
                             ).str() );
    }

  };

  //-------------------------------------------------------------------------------
  // pcps::MoveFunc -- Specialization for spins
  //-------------------------------------------------------------------------------

  template <class SCALAR, class REF> class MoveFunc<SCALAR, pcps::Part_Spin, REF> {

    // private data
    private:
      const int _configs_per_site;
      const int _nsites;
      const int _target_occ;
      const bool _project_sz;

    public:

    // constructor
    MoveFunc(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, const pcps::Wavefunction<SCALAR> & wfn, pcps::Arena & arena)
      : _configs_per_site(userinp.configs_per_site()),
        _nsites(userinp.nsites()),
        _target_occ(0.5 * userinp.nsites() * ( userinp.configs_per_site() - 1 ) + userinp.z_spin() + 0.01),
        _project_sz(userinp.project_sz())
    {
      if ( std::abs( double(_target_occ) - ( 0.5 * userinp.nsites() * ( userinp.configs_per_site() - 1 ) + userinp.z_spin() ) ) > 1.0e-6 )
        throw pcps::Exception( (boost::format("error in MoveFunc::init():  target occupation is incorrect")).str() );
    }

    // function to set the initial configuration
    template<class GENERATOR> void initialize_configuration(const pcps::Input & userinp,
                                                            pcps::CoeffValueFunc<SCALAR, pcps::Part_Spin, REF> & cvf,
                                                            int * const config,
                                                            GENERATOR & gen)
    {

      // if supplied, start in the user's chosen configuration
      if (userinp.init_config().size() == _nsites) {

        for (int i = 0; i < _nsites; i++)
          config[i] = userinp.init_config().at(i);

      // otherwise, generate random configurations until we get one with an acceptable wavefunction value
      } else {

        do {

          // random config with specific sz value
          if (_project_sz) {

            for (int i = 0; i < _nsites; i++)
              config[i] = 0;
            for (int occ_remaining = _target_occ; occ_remaining > 0; ) {
              const int site = gen() % _nsites;
              if (config[site] < _configs_per_site-1) {
                config[site]++;
                occ_remaining--;
              }
            }

          // random config
          } else {

            for (int i = 0; i < _nsites; i++)
              config[i] = gen() % _configs_per_site;

          }

        } while ( ! pcps::good_init_config_value(userinp, cvf, config) );

      }

    }

    // function to propose and make a move and return the ratio of proposal densities for the configs
    template <class GENERATOR> double propose(int & n_affected, int * const affected, int * const config, GENERATOR & gen) {

      // get the maximum and minimum occupations for a site
      const int max_occ = _configs_per_site - 1;
      const int min_occ = 0;

      for (int i = 0; i < _nsites; i++) {
        assert( config[i] <= max_occ );
        assert( config[i] >= min_occ );
      }

      if (_project_sz) {

        // get the position of the sites to modify and record their old occupations
        int sum_occ, pos1, pos2, pos1_old_occ, pos2_old_occ;
        do {
          pos1 = gen() % _nsites;
          pos2 = gen() % (_nsites-1);
          if (pos2 >= pos1)
            pos2++;
          pos1_old_occ = config[pos1];
          pos2_old_occ = config[pos2];
          sum_occ = pos1_old_occ + pos2_old_occ;
        } while ( sum_occ == 2*min_occ || sum_occ == 2*max_occ );

        // determine the new maximum and minimum occupations possible for each site
        const int max_new_occ = std::min(max_occ, sum_occ);
        const int min_new_occ = std::max(min_occ, sum_occ - max_occ);

        // select a new occupation that preserves the sites' total occupation
        config[pos1] = min_new_occ + gen() % ( max_new_occ - min_new_occ );
        if (config[pos1] >= pos1_old_occ)
          config[pos1]++;
        config[pos2] = sum_occ - config[pos1];
        //std::cout << boost::format("%i %i %i") % sum_occ % config[pos1] % config[pos2] << std::endl;

        // record how many and which sites were affected by the move
        n_affected = 2;
        affected[0] = pos1;
        affected[1] = pos2;

      } else {

        const int pos = gen() % _nsites;
        int occ = gen() % max_occ;
        if (occ == config[pos])
          occ++;
        config[pos] = occ;
        n_affected = 1;
        affected[0] = pos;

      }

      for (int i = 0; i < _nsites; i++) {
        assert( config[i] <= max_occ );
        assert( config[i] >= min_occ );
      }

      return 1.0;

    }

    // function to update internal data after a move has been accepted
    void accept(const int n_affected, const int * const affected, const int * const new_config) {}

  };

  //-------------------------------------------------------------------------------
  // pcps::MoveFunc -- Specialization for fermions
  //-------------------------------------------------------------------------------

  template <class SCALAR, class REF> class MoveFunc<SCALAR, pcps::Part_Fermion, REF> {

    // private data
    private:
      const int _nsites;
      const int _n_a_sites;
      const int _n_total_particles;
      const double _z_spin;
      const bool _exchange_move;
      const bool _mix_move;
      const bool _double_ab_move; ///< whether to use a combined alpha/beta move in addition to the standard move
      int _total_partners; // number of valid site pairs for mixing move
      int * _n_partners; // number of valid partners for each site for mixing move

    public:

    // constructor
    MoveFunc(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, const pcps::Wavefunction<SCALAR> & wfn, pcps::Arena & arena)
      : _nsites(userinp.nsites()),
        _n_a_sites(userinp.nsites()/2),
        _n_total_particles(userinp.nparticles()),
        _z_spin(userinp.z_spin()),
        _exchange_move(userinp.move_type() == pcps::keyword::Move_Exchange),
        _mix_move(userinp.move_type() == pcps::keyword::Move_Mix),
        _double_ab_move(userinp.move_type() == pcps::keyword::Move_DoubleAB),
        _n_partners(NULL)
    {
      _n_partners = arena.allocate_array<     int>( size_t(_n_a_sites), 0);
    }

    // function to get the number of particles, 0 = alpha, 1 = beta
    int n_particles(const int b) const { return pcps::round( 0.5 * _n_total_particles + (1-2*b) * _z_spin ); }

    // function to set the initial configuration
    template<class GENERATOR> void initialize_configuration(const pcps::Input & userinp,
                                                            pcps::CoeffValueFunc<SCALAR, pcps::Part_Fermion, REF> & cvf,
                                                            int * const config,
                                                            GENERATOR & gen)
    {

      assert(_nsites == 2 * _n_a_sites);
      assert(this->n_particles(0) <= _n_a_sites);
      assert(this->n_particles(1) <= _n_a_sites);

      // if supplied, start in the user's chosen configuration
      if (userinp.init_config().size() == _nsites) {

        for (int i = 0; i < _nsites; i++)
          config[i] = userinp.init_config().at(i);

      // otherwise, generate random configurations until we get one with an acceptable wavefunction value
      } else {

        do {

          // remove all particles from the configuration
          for (int i = 0; i < _nsites; i++)
            config[i] = 0;

          // place alpha (b == 0) and beta (b == 1) particles
          for (int b = 0; b <= 1; b++) {
            for (int i = 0; i < this->n_particles(b); i++) {
              int x = gen() % (_n_a_sites - i);
              int c = 0;
              int j = 0;
              for ( ; j < _n_a_sites; j++)
                if (config[2*j+b] == 0)
                  if (c++ == x)
                    break;
              assert(c == x+1);
              assert(config[2*j+b] == 0);
              config[2*j+b] = 1;
            }
          }

        } while ( ! pcps::good_init_config_value(userinp, cvf, config) );

      }

      // if we are using the mixing move, initialize the list of how many valid mixing partners each site has and count the total number of valid pairs
      if (_mix_move) {
        _total_partners = 0;
        for (int i = 0; i < _n_a_sites; i++) {
          _n_partners[i] = 0;
          for (int j = 0; j < _n_a_sites; j++)
            if ( config[2*i] != config[2*j] || config[2*i+1] != config[2*j+1] )
              _n_partners[i]++;
          _total_partners += _n_partners[i];
        }
        assert( _total_partners % 2 == 0 );
        _total_partners /= 2; // compensate for the fact that we have counted each valid pair twice
      }

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   Chooses a particle of the given type at random and return its position.
    ///
    /// \param[in]     config   the current lattice configuration
    /// \param[in]     b        tye type of particle to choose (0 = alpha, 1 = beta)
    /// \param[in,out] gen      a random number generator
    ///
    /// \return The position of the randomly chosen particle.
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    template <class GENERATOR> int get_rand_part_pos(const int * const config, const int b, GENERATOR & gen) {
      int x = gen() % this->n_particles(b);
      for (int i = 0; i < _nsites; i+=2)
        if (config[i+b] == 1)
          if (x-- == 0)
            return i+b;
      throw pcps::Exception("failed to choose a particle in pcps::MoveFunc::get_rand_part_pos");
      return -1;
    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   Chooses a hole of the given type at random and return its position.
    ///
    /// \param[in]     config   the current lattice configuration
    /// \param[in]     b        tye type of hole to choose (0 = alpha, 1 = beta)
    /// \param[in,out] gen      a random number generator
    ///
    /// \return The position of the randomly chosen hole.
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    template <class GENERATOR> int get_rand_hole_pos(const int * const config, const int b, GENERATOR & gen) {
      int x = gen() % (_n_a_sites - this->n_particles(b));
      for (int i = 0; i < _nsites; i+=2)
        if (config[i+b] == 0)
          if (x-- == 0)
            return i+b;
      throw pcps::Exception("failed to choose a hole in pcps::MoveFunc::get_rand_hole_pos");
      return -1;
    }

    // function to propose and make a move and return the ratio of proposal densities for the configs
    template <class GENERATOR> double propose(int & n_affected, int * const affected, int * const config, GENERATOR & gen) {

      // move that exchanges the occupations of two sites (assumes single occupancy and opposite spin)
      if (_exchange_move) {

        // get the position of the alpha and beta sites to modify
        int pos_a, pos_b;
        do { pos_a = 2 * ( gen() % _n_a_sites )    ; } while (config[pos_a] == 0                    );
        do { pos_b = 2 * ( gen() % _n_a_sites ) + 1; } while (config[pos_b] == 0 || pos_a == pos_b-1);

        assert( config[pos_a+1] == 0 );
        assert( config[pos_b-1] == 0 );

        // make the move
        config[pos_a  ] = 0; // alpha move
        config[pos_b-1] = 1; // alpha move
        config[pos_b  ] = 0; //  beta move
        config[pos_a+1] = 1; //  beta move
        n_affected = 4;
        affected[0] = pos_a  ; // alpha move
        affected[1] = pos_b-1; // alpha move
        affected[2] = pos_b  ; //  beta move
        affected[3] = pos_a+1; //  beta move

        return 1.0;

      }

      // move that mixes the occupations of two sites
      if (_mix_move) {

        // get the positions of the sites to mix
        int pos_a1, pos_a2, pos_b1, pos_b2;
        do {
          pos_a1 = 2 * ( gen() % _n_a_sites );
          pos_a2 = 2 * ( gen() % (_n_a_sites-1) );
          if (pos_a2 >= pos_a1)
            pos_a2 += 2;
          pos_b1 = pos_a1 + 1;
          pos_b2 = pos_a2 + 1;
        } while ( config[pos_a1] == config[pos_a2] && config[pos_b1] == config[pos_b2] );

        // check sanity of occupations
        assert( config[pos_a1] == 0 || config[pos_a1] == 1 );
        assert( config[pos_a2] == 0 || config[pos_a2] == 1 );
        assert( config[pos_b1] == 0 || config[pos_b1] == 1 );
        assert( config[pos_b2] == 0 || config[pos_b2] == 1 );

        // function object to swap electrons
        struct Local_Swap {
          void operator() (const int pos1, const int pos2, int * const config, int * const affected) {
            affected[0] = ( config[pos1] == 0 ? pos1 : pos2 );
            affected[1] = ( config[pos1] == 0 ? pos2 : pos1 );
            std::swap(config[pos1], config[pos2]);
          }
        };
        Local_Swap electron_swap;

        // initialize the proposal density ratio
        double q = 1.0;

        n_affected = 2;

        // sites differ in both alpha and beta occupations
        if ( config[pos_a1] != config[pos_a2] && config[pos_b1] != config[pos_b2] ) {

          const int x = gen() % 10;  // value of x determines the move as follows:  0 - change alpha   1 - change beta   2+ - change both
          if (x == 0) {
            electron_swap(pos_a1, pos_a2, config, affected);
          } else if (x == 1) {
            electron_swap(pos_b1, pos_b2, config, affected);
          } else if (x >= 2) {
            electron_swap(pos_a1, pos_a2, config, affected+0);
            electron_swap(pos_b1, pos_b2, config, affected+2);
            n_affected = 4;
          } else throw pcps::Exception( (boost::format("unexpected x value of %i when making a mixed move") % x).str() );

          // get the ratio of the reverse proposal density to the forward proposal density
          int new_total_partners = _total_partners - _n_partners[pos_a1/2] - _n_partners[pos_a2/2];
          for (int i = 0; i < _n_a_sites; i++) {
            if ( config[2*i] != config[pos_a1] || config[2*i+1] != config[pos_b1] ) new_total_partners++;
            if ( config[2*i] != config[pos_a2] || config[2*i+1] != config[pos_b2] ) new_total_partners++;
          }
          q = double(_total_partners) / double(new_total_partners);

        // sites differ in alpha occupation only
        } else if (config[pos_a1] != config[pos_a2]) {

          electron_swap(pos_a1, pos_a2, config, affected);

        // sites differ in alpha occupation only
        } else if (config[pos_b1] != config[pos_b2]) {

          electron_swap(pos_b1, pos_b2, config, affected);

        // sites do not differ
        } else throw pcps::Exception( (boost::format("site occupations must differ when attempting to make a mixed move")).str() );

        return q;

      }

      // double move, which moves 1 alpha electron and 1 beta electron
      if ( _double_ab_move && pcps::uni_01(gen) >= 0.5 ) {

        // choose alpha and beta electrons at random and get their positions
        const int ap = this->get_rand_part_pos(config, 0, gen);
        const int bp = this->get_rand_part_pos(config, 1, gen);

        // choose alpha and beta holes at random and get their positions
        const int ah = this->get_rand_hole_pos(config, 0, gen);
        const int bh = this->get_rand_hole_pos(config, 1, gen);

        // exchange particles and holes
        std::swap(config[ap], config[ah]);
        std::swap(config[bp], config[bh]);

        // record how many and which sites were affected
        n_affected = 4;
        affected[0] = ah;
        affected[1] = ap;
        affected[2] = bh;
        affected[3] = bp;

        return 1.0;

      }

      // standard move, which moves 1 alpha electron or 1 beta electron

      // get the position of the sites to modify and swap their occupations
      const int pos1 = gen() % _nsites;
      const int pos1_occ = config[pos1];
      assert(pos1_occ == 0 || pos1_occ == 1);
      const int b = pos1 % 2;
      const int x = gen() % ( pos1_occ == 1 ? _n_a_sites - this->n_particles(b) : this->n_particles(b) );
      int c = 0;
      int j = 0;
      for ( ; j < _n_a_sites; j++)
        if (config[2*j+b] != pos1_occ)
          if (c++ == x)
            break;
      const int pos2 = 2*j+b;
      assert(c == x+1);
      assert(pos1 != pos2);
      assert(pos1 % 2 == pos2 % 2);
      assert(config[pos1] != config[pos2]);
      std::swap(config[pos1], config[pos2]);

      // record how many and which sites were affected by the move
      n_affected = 2;
      affected[0] = pos1;
      affected[1] = pos2;

      return 1.0;

    }

    // function to update internal data after a move has been accepted
    void accept(const int n_affected, const int * const affected, const int * const new_config) {

      // update the information on valid mixing partners
      if (_mix_move) {
        const int pos1 = affected[0]/2;
        const int pos2 = affected[1]/2;
        assert( pos1 != pos2 );
        _total_partners -= _n_partners[pos1];
        _total_partners -= _n_partners[pos2];
        _n_partners[pos1] = 0;
        _n_partners[pos2] = 0;
        for (int i = 0; i < _n_a_sites; i++) {
          if ( new_config[2*i] != new_config[2*pos1] || new_config[2*i+1] != new_config[2*pos1+1] ) _n_partners[pos1]++;
          if ( new_config[2*i] != new_config[2*pos2] || new_config[2*i+1] != new_config[2*pos2+1] ) _n_partners[pos2]++;
        }
        _total_partners += _n_partners[pos1];
        _total_partners += _n_partners[pos2];
      }

    }

  };

} // end namespace pcps

#endif
